import os
from transformers import AutoTokenizer, AutoModelForMaskedLM


model_path = r'C:\Users\admin\.cache\huggingface\hub\models--bert-base-chinese\snapshots\8d2a91f91cc38c96bb8b4556ba70c392f8d5ee55'


assert os.path.exists(model_path), '获取预训练模型失败！`model_path`不存在？'


def get_tokenizer():
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    # pretrained_model = AutoModelForMaskedLM.from_pretrained(model_path)
    return tokenizer


def get_pretrained_model__bert():
    from transformers import BertModel
    pretrained_model = BertModel.from_pretrained(model_path)
    # pretrained_model = AutoModelForMaskedLM.from_pretrained(model_path)
    return pretrained_model


import time


class WithTimer:
    """
    # 使用示例
    with WithTimer("test_func"):
        # 在这里放置要测量执行时间的代码
        for _ in range(1000):
            pass
    """
    mute_all = False
    mute_enter = False
    mute_exit = False

    def __init__(self, name, tt=None, debug=None):
        self.name = name

        if tt is not None:
            assert hasattr(tt, 'now'), 'tt必须实现方法`now`, 以获取当前总花费时间!'
        self.tt = tt

        if debug is None:
            debug = not WithTimer.mute_all
        self.debug = debug

    def __enter__(self):
        if self.debug and not WithTimer.mute_enter:
            msg = f"====== enter WithTimer[{self.name}]"
            if self.tt:
                msg += f", now: {self.tt.now()}"

            print()
            print(msg)
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        elapsed_time = self.end_time - self.start_time
        if self.debug and not WithTimer.mute_exit:
            msg = f"====== exit WithTimer[{self.name}] cost seconds: {elapsed_time:.3f}"
            if self.tt:
                msg += f', now: {self.tt.now()}'
            print(msg)


with_timer = WithTimer


if __name__ == '__main__':
    tokenizer = get_tokenizer()
    pretrained_model = get_pretrained_model__bert()









